{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "from pyspark import SparkContext\n",
    "from pyspark.sql import SparkSession\n",
    "from pyspark.ml.classification import DecisionTreeClassificationModel,DecisionTreeClassifier\n",
    "import pyspark.ml.linalg as linalg\n",
    "from pyspark.ml.feature import StringIndexer,VectorAssembler\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "spark = SparkSession.builder.getOrCreate()\n",
    "sc = spark.sparkContext"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 107,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 生成测试数据集\n",
    "data = {\n",
    "         'Age': {0: 22, 1: 38, 2: 26, 3: 35, 4: 35, 5: 40},\n",
    "         'Fare': {0: 7.3, 1: 71.3, 2: 7.9, 3: 53.1, 4: 8.0, 5: 53.1},\n",
    "         'Pclass': {0: 3, 1: 1, 2: 3, 3: 1, 4: 3, 5: 1},\n",
    "         'Survived': {0: 0., 1: 1., 2: 1., 3: 1., 4: 0.,5: 1.}}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {},
   "outputs": [],
   "source": [
    "#转换pandas df到spark df\n",
    "df_pd = pd.DataFrame(data)\n",
    "df_spark = spark.createDataFrame(df_pd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---+----+------+--------+\n",
      "|Age|Fare|Pclass|Survived|\n",
      "+---+----+------+--------+\n",
      "| 22| 7.3|     3|     0.0|\n",
      "| 38|71.3|     1|     1.0|\n",
      "| 26| 7.9|     3|     1.0|\n",
      "| 35|53.1|     1|     1.0|\n",
      "| 35| 8.0|     3|     0.0|\n",
      "| 40|53.1|     1|     1.0|\n",
      "+---+----+------+--------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df_spark.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 数据预处理（最简单的特征工程）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---+----+------+--------+-------------+\n",
      "|Age|Fare|Pclass|Survived|indexed_label|\n",
      "+---+----+------+--------+-------------+\n",
      "| 22| 7.3|     3|     0.0|          1.0|\n",
      "| 38|71.3|     1|     1.0|          0.0|\n",
      "| 26| 7.9|     3|     1.0|          0.0|\n",
      "| 35|53.1|     1|     1.0|          0.0|\n",
      "| 35| 8.0|     3|     0.0|          1.0|\n",
      "| 40|53.1|     1|     1.0|          0.0|\n",
      "+---+----+------+--------+-------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "indexer = StringIndexer(inputCol=\"Survived\", outputCol=\"indexed_label\").fit(df_spark)\n",
    "indexed_df = indexer.transform(df_spark)\n",
    "indexed_df.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {},
   "outputs": [],
   "source": [
    "assembler = VectorAssembler(inputCols=['Age','Fare','Pclass'],outputCol='features')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---+----+------+--------+-------------+---------------+\n",
      "|Age|Fare|Pclass|Survived|indexed_label|       features|\n",
      "+---+----+------+--------+-------------+---------------+\n",
      "| 22| 7.3|     3|     0.0|          1.0| [22.0,7.3,3.0]|\n",
      "| 38|71.3|     1|     1.0|          0.0|[38.0,71.3,1.0]|\n",
      "| 26| 7.9|     3|     1.0|          0.0| [26.0,7.9,3.0]|\n",
      "| 35|53.1|     1|     1.0|          0.0|[35.0,53.1,1.0]|\n",
      "| 35| 8.0|     3|     0.0|          1.0| [35.0,8.0,3.0]|\n",
      "| 40|53.1|     1|     1.0|          0.0|[40.0,53.1,1.0]|\n",
      "+---+----+------+--------+-------------+---------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "final_df = assembler.transform(indexed_df)\n",
    "final_df.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 120,
   "metadata": {},
   "outputs": [],
   "source": [
    "#拆分数据集\n",
    "train, test = final_df.randomSplit([0.8,0.2])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 模型训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 122,
   "metadata": {},
   "outputs": [],
   "source": [
    "dt = DecisionTreeClassifier(maxDepth=2, labelCol=\"indexed_label\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 123,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = dt.fit(final_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 127,
   "metadata": {},
   "outputs": [],
   "source": [
    "#预测\n",
    "result = model.transform(test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 134,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---+----+------+--------+-------------+--------------+-------------+-----------+----------+\n",
      "|Age|Fare|Pclass|Survived|indexed_label|      features|rawPrediction|probability|prediction|\n",
      "+---+----+------+--------+-------------+--------------+-------------+-----------+----------+\n",
      "| 26| 7.9|     3|     1.0|          0.0|[26.0,7.9,3.0]|    [1.0,1.0]|  [0.5,0.5]|       0.0|\n",
      "+---+----+------+--------+-------------+--------------+-------------+-----------+----------+\n",
      "\n",
      "+---+----+------+--------+-------------+--------------+\n",
      "|Age|Fare|Pclass|Survived|indexed_label|      features|\n",
      "+---+----+------+--------+-------------+--------------+\n",
      "| 26| 7.9|     3|     1.0|          0.0|[26.0,7.9,3.0]|\n",
      "+---+----+------+--------+-------------+--------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "#对比预测与实际\n",
    "result.show()\n",
    "test.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 125,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DecisionTreeClassificationModel (uid=DecisionTreeClassifier_442e36057159) of depth 2 with 5 nodes"
      ]
     },
     "execution_count": 125,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.save()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 文档示例"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = spark.createDataFrame([\n",
    "  (1.0, Vectors.dense(1.0,1)),\n",
    "     (0.0, Vectors.dense(1,0))], [\"label\", \"features\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-----+---------+\n",
      "|label| features|\n",
      "+-----+---------+\n",
      "|  1.0|[1.0,1.0]|\n",
      "|  0.0|[1.0,0.0]|\n",
      "+-----+---------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "metadata": {},
   "outputs": [],
   "source": [
    "stringIndexer = StringIndexer(inputCol=\"label\", outputCol=\"indexed\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {},
   "outputs": [],
   "source": [
    "si_model = stringIndexer.fit(df)\n",
    "td = si_model.transform(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-----+---------+-------+\n",
      "|label| features|indexed|\n",
      "+-----+---------+-------+\n",
      "|  1.0|[1.0,1.0]|    1.0|\n",
      "|  0.0|[1.0,0.0]|    0.0|\n",
      "+-----+---------+-------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "td.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
